#ifndef THREAD_POOL_H
#define THREAD_POOL_H

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>

/* 平台检测与头文件包含 */
#ifdef _WIN32
    #include <windows.h>
    #include <synchapi.h>
#else
    #include <pthread.h>
    #include <semaphore.h>
#endif

/* 类型定义 */
typedef void* (*routine_t)(void*);  // 任务函数指针类型

// 任务结构体
typedef struct {
    routine_t routine;  // 任务函数
    void* args;         // 任务参数
} task_t;

// 队列节点结构体
typedef struct queue_node {
    void* data;
    struct queue_node* next;
} queue_node_t;

// 队列结构体
typedef struct {
    queue_node_t* head;
    queue_node_t* tail;
    size_t size;
} queue_t;

// 跨平台类型定义
#ifdef _WIN32
    typedef HANDLE thread_t;          // Windows线程句柄
    typedef CRITICAL_SECTION mutex_t; // Windows临界区(互斥锁)
    typedef HANDLE semaphore_t;       // Windows信号量
#else
    typedef pthread_t thread_t;       // POSIX线程
    typedef pthread_mutex_t mutex_t;  // POSIX互斥锁
    typedef sem_t semaphore_t;        // POSIX信号量
#endif

// 线程池结构体
typedef struct {
    uint8_t size;               // 线程数量
    thread_t* threads;          // 线程数组
    queue_t queue;              // 任务队列
    mutex_t tpool_lock;         // 互斥锁
    semaphore_t empty;          // 信号量，用于等待任务
} tpool_t;

/* 颜色定义，用于信息打印 */
typedef enum {
    COLOR_RED,
    COLOR_GREEN,
    COLOR_YELLOW,
    COLOR_BLUE,
    COLOR_MAGENTA,
    COLOR_CYAN,
    COLOR_WHITE
} color_t;

/* 终端颜色控制宏 */
#define RESET       "\033[0m"
#define BOLD        1
#define NO_BOLD     0

/* 函数声明 */

// 队列操作函数
void queue_new(queue_t* q);
void queue_push(queue_t* q, void* data);
void* queue_pop(queue_t* q);
void queue_free(queue_t* q);

// 线程池操作函数
void tpool_new(tpool_t* pool, uint8_t size);
void tpool_push(tpool_t* pool, routine_t routine, void* args);
task_t* tpool_pop(tpool_t* pool);
void* tpool_stop(void* args);
void tpool_finish(tpool_t* pool);
void* task_worker(void* args);

// 工具函数
void col_str_style(color_t color, int bold, char* buf);
void print_info(const char* msg, color_t color, bool exit_process,
                bool print_line, const char* file, int line);

// 错误处理宏
#define error(msg) print_info(msg, COLOR_RED, true, true, __FILE__, __LINE__)

#endif // THREAD_POOL_H

#ifdef THREAD_POOL_IMPLEMENTATION

/* 队列实现 */
void queue_new(queue_t* q) {
    q->head = NULL;
    q->tail = NULL;
    q->size = 0;
}

void queue_push(queue_t* q, void* data) {
    queue_node_t* node = (queue_node_t*)malloc(sizeof(queue_node_t));
    if (!node) {
        error("内存分配失败");
    }
    node->data = data;
    node->next = NULL;

    if (q->tail == NULL) {
        q->head = node;
        q->tail = node;
    } else {
        q->tail->next = node;
        q->tail = node;
    }
    q->size++;
}

void* queue_pop(queue_t* q) {
    if (q->head == NULL) return NULL;

    queue_node_t* node = q->head;
    void* data = node->data;
    q->head = q->head->next;

    if (q->head == NULL) {
        q->tail = NULL;
    }

    free(node);
    q->size--;
    return data;
}

void queue_free(queue_t* q) {
    while (q->head != NULL) {
        queue_node_t* node = q->head;
        q->head = q->head->next;
        free(node->data);  // 释放节点中存储的数据
        free(node);        // 释放节点本身
    }
    q->tail = NULL;
    q->size = 0;
}

/* 跨平台线程函数适配器 */
#ifdef _WIN32
// Windows线程函数适配器，转换为pthread风格
DWORD WINAPI thread_adapter(LPVOID args) {
    // 调用实际的工作函数
    task_worker(args);
    return 0;
}
#endif

/* 线程池工作函数实现 */
void* task_worker(void* args) {
    tpool_t* pool = (tpool_t*)args;

    while (1) {
        // 从线程池获取任务，会阻塞直到有任务可用
        task_t* popped_task = tpool_pop(pool);

        // 如果是停止任务，则退出循环
        if (popped_task->routine == tpool_stop) {
            free(popped_task);  // 释放停止任务
            break;
        }

        // 执行任务
        popped_task->routine(popped_task->args);
        free(popped_task);  // 释放任务
    }
    return NULL;
}

/* 线程池实现 */
void tpool_new(tpool_t* pool, uint8_t size) {
    pool->size = size;
    pool->threads = (thread_t*)malloc(size * sizeof(thread_t));
    if (!pool->threads) {
        error("内存分配失败");
    }

    // 初始化任务队列
    queue_new(&pool->queue);

    // 初始化互斥锁
#ifdef _WIN32
    InitializeCriticalSection(&pool->tpool_lock);
#else
    if (pthread_mutex_init(&pool->tpool_lock, NULL) != 0) {
        error("初始化互斥锁失败");
    }
#endif

    // 初始化信号量，用于等待任务
#ifdef _WIN32
    // Windows: 初始值0，最大计数为SIZE_MAX
   // Windows平台使用LONG_MAX作为信号量最大值
    pool->empty = CreateSemaphore(NULL, 0, LONG_MAX, NULL);
    if (pool->empty == NULL) {
        error("初始化信号量失败");
    }
#else
    if (sem_init(&pool->empty, 0, 0) != 0) {
        error("初始化信号量失败");
    }
#endif

    // 创建工作线程
#ifdef _WIN32
    EnterCriticalSection(&pool->tpool_lock);
#else
    pthread_mutex_lock(&pool->tpool_lock);
#endif

    for (uint8_t i = 0; i < size; i++) {
#ifdef _WIN32
        pool->threads[i] = CreateThread(
            NULL,           // 默认安全属性
            0,              // 默认栈大小
            thread_adapter, // 线程函数(适配器)
            pool,           // 传递给线程函数的参数
            0,              // 立即执行线程
            NULL            // 不需要线程ID
        );
        if (pool->threads[i] == NULL) {
            error("创建线程失败");
        }
#else
        if (pthread_create(&pool->threads[i], NULL, task_worker, (void*)pool) != 0) {
            error("创建线程失败");
        }
#endif
    }

#ifdef _WIN32
    LeaveCriticalSection(&pool->tpool_lock);
#else
    pthread_mutex_unlock(&pool->tpool_lock);
#endif
}

/**
 * @brief 一个空函数，当线程执行它时，会停止并销毁自身
 */
void* tpool_stop(void* args) {
    // 空实现，仅作为停止标记
    return NULL;
}

/**
 * @brief 向线程池添加任务
 *
 * @param pool 线程池实例
 * @param routine 任务函数
 * @param args 任务参数
 */
void tpool_push(tpool_t* pool, routine_t routine, void* args) {
    // 创建新任务
    task_t* new_task = (task_t*)malloc(sizeof(task_t));
    if (!new_task) {
        error("内存分配失败");
    }
    new_task->routine = routine;
    new_task->args = args;

    // 加锁并将任务加入队列
#ifdef _WIN32
    EnterCriticalSection(&pool->tpool_lock);
#else
    pthread_mutex_lock(&pool->tpool_lock);
#endif

    queue_push(&pool->queue, (void*)new_task);
    
#ifdef _WIN32
    ReleaseSemaphore(pool->empty, 1, NULL);  // 发送信号，通知有新任务
#else
    sem_post(&pool->empty);  // 发送信号，通知有新任务
#endif

#ifdef _WIN32
    LeaveCriticalSection(&pool->tpool_lock);
#else
    pthread_mutex_unlock(&pool->tpool_lock);
#endif
}

/**
 * @brief 从线程池获取任务
 *
 * @param pool 线程池实例
 * @return task_t* 获取到的任务
 */
task_t* tpool_pop(tpool_t* pool) {
    // 等待直到有任务可用
#ifdef _WIN32
    WaitForSingleObject(pool->empty, INFINITE);
#else
    sem_wait(&pool->empty);
#endif

    task_t* popped_task;

    // 加锁并从队列中取出任务
#ifdef _WIN32
    EnterCriticalSection(&pool->tpool_lock);
#else
    pthread_mutex_lock(&pool->tpool_lock);
#endif

    popped_task = (task_t*)queue_pop(&pool->queue);

#ifdef _WIN32
    LeaveCriticalSection(&pool->tpool_lock);
#else
    pthread_mutex_unlock(&pool->tpool_lock);
#endif

    return popped_task;
}

/**
 * @brief 销毁线程池，等待所有线程完成并释放资源
 *
 * @param pool 要销毁的线程池
 */
void tpool_finish(tpool_t* pool) {
    // 向线程池添加停止任务，每个线程一个
    for (uint8_t i = 0; i < pool->size; i++) {
        tpool_push(pool, tpool_stop, NULL);
    }

    // 等待所有线程结束
    for (uint8_t i = 0; i < pool->size; i++) {
#ifdef _WIN32
        WaitForSingleObject(pool->threads[i], INFINITE);
        CloseHandle(pool->threads[i]);
#else
        pthread_join(pool->threads[i], NULL);
#endif
    }

    // 释放资源
    free(pool->threads);
    
#ifdef _WIN32
    DeleteCriticalSection(&pool->tpool_lock);
    CloseHandle(pool->empty);
#else
    pthread_mutex_destroy(&pool->tpool_lock);
    sem_destroy(&pool->empty);
#endif
    
    queue_free(&pool->queue);
}

/* 工具函数实现 */
/**
 * @brief 设置终端输出颜色和样式
 *
 * @param color 颜色
 * @param bold 是否粗体
 * @param buf 存储颜色控制字符串的缓冲区
 */
void col_str_style(color_t color, int bold, char* buf) {
    const char* color_code;
    switch (color) {
        case COLOR_RED:     color_code = "31"; break;
        case COLOR_GREEN:   color_code = "32"; break;
        case COLOR_YELLOW:  color_code = "33"; break;
        case COLOR_BLUE:    color_code = "34"; break;
        case COLOR_MAGENTA: color_code = "35"; break;
        case COLOR_CYAN:    color_code = "36"; break;
        case COLOR_WHITE:   color_code = "37"; break;
        default:            color_code = "0";  break;
    }

    if (bold) {
        sprintf(buf, "\033[%s;1m", color_code);
    } else {
        sprintf(buf, "\033[%sm", color_code);
    }
}

/**
 * @brief 打印信息，支持颜色和样式
 *
 * @param msg 要打印的消息
 * @param color 消息颜色
 * @param exit_process 是否在打印后退出程序
 * @param print_line 是否打印行号和文件名
 * @param file 文件名
 * @param line 行号
 */
void print_info(const char* msg, color_t color, bool exit_process,
                bool print_line, const char* file, int line) {
    char text_col_bold[10] = {0};
    col_str_style(color, BOLD, text_col_bold);

    if (print_line) {
        printf("\n%s%s 行 %d: %s%s\n",
               text_col_bold,
               file,
               line,
               msg,
               RESET);
    } else {
        printf("%s%s%s\n", text_col_bold, msg, RESET);
    }

    if (exit_process) {
        exit(1);
    }
}

#endif // THREAD_POOL_IMPLEMENTATION
